#include <cuda_runtime.h>

#include "util.h"

static void perform_copy_cutting_internal(int rank, int typeSize, long *header1, long *header2, Interval *cutting, enum cudaMemcpyKind where, int matched,
        int skip, long offs1, long offs2) {
    int effectiveRank = rank - skip;
    while (effectiveRank > 1 && cutting[rank - effectiveRank][0] == cutting[rank - effectiveRank][1])
        effectiveRank--;
    dvmh_log(TRACE, "effective rank = %d", effectiveRank);
    if (matched >= effectiveRank - 1) {
        // YAHOO! We can copy it by one piece
        char *address1 = (char *)header1[rank + 2] + typeSize * (header1[rank + 1] + offs1);
        char *address2 = (char *)header2[rank + 2] + typeSize * (header2[rank + 1] + offs2);
        long elemCount = 1;
        int i;
        for (i = skip; i < rank - 1; i++) {
            address1 += typeSize * header1[i + 1] * cutting[i][0];
            address2 += typeSize * header2[i + 1] * cutting[i][0];
            elemCount *= cutting[i][1] - cutting[i][0] + 1;
        }
        if (rank > 0) {
            address1 += typeSize * cutting[rank - 1][0];
            address2 += typeSize * cutting[rank - 1][0];
            elemCount *= cutting[rank - 1][1] - cutting[rank - 1][0] + 1;
        }

        dvmh_log(TRACE, "copying %ld bytes (%ld elements)", elemCount * typeSize, elemCount);
        assertCuda(cudaMemcpy(address2, address1, elemCount * typeSize, where));
    } else if (matched == rank - skip - 2) {
        // YAHOO! We can copy it by one 2D command
        char *address1 = (char *)header1[rank + 2] + typeSize * (header1[rank + 1] + offs1);
        char *address2 = (char *)header2[rank + 2] + typeSize * (header2[rank + 1] + offs2);
        int i;
        for (i = skip; i < rank - 1; i++) {
            address1 += typeSize * header1[i + 1] * cutting[i][0];
            address2 += typeSize * header2[i + 1] * cutting[i][0];
        }
        address1 += typeSize * cutting[rank - 1][0];
        address2 += typeSize * cutting[rank - 1][0];

        long height = cutting[skip][1] - cutting[skip][0] + 1;
        long width = typeSize;
        long pitch1 = header1[skip + 1] * typeSize;
        long pitch2 = header2[skip + 1] * typeSize;
        for (i = skip + 1; i < rank; i++)
            width *= cutting[i][1] - cutting[i][0] + 1;

        dvmh_log(TRACE, "offs1=%ld offs2=%ld address1=%p address2=%p where=%d", offs1, offs2, address1, address2, where);
        dvmh_log(TRACE, "copying %ld bytes (2D, pitch1=%ld pitch2=%ld width=%ld height=%ld)", height * width, pitch1, pitch2, width, height);
        assertCuda(cudaMemcpy2D(address2, pitch2, address1, pitch1, width, height, where));
    } else {
        // Oh no, we can not copy it by one command... :(
        long curVal;
        for (curVal = cutting[skip][0]; curVal <= cutting[skip][1]; curVal++) {
            long offset1 = offs1 + header1[skip + 1] * curVal;
            long offset2 = offs2 + header2[skip + 1] * curVal;
            perform_copy_cutting_internal(rank, typeSize, header1, header2, cutting, where, matched, skip + 1, offset1, offset2);
        }
    }
}

static void perform_copy_peer_internal(int rank, int typeSize, int device1, long *header1, int device2, long *header2, Interval *cutting, int matched,
        int skip, long offs1, long offs2) {
    if (matched >= rank - skip - 1) {
        // YAHOO! We can copy it by one piece
        char *address1 = (char *)header1[rank + 2] + typeSize * (header1[rank + 1] + offs1);
        char *address2 = (char *)header2[rank + 2] + typeSize * (header2[rank + 1] + offs2);
        long elemCount = 1;
        int i;
        for (i = skip; i < rank - 1; i++) {
            address1 += typeSize * header1[i + 1] * cutting[i][0];
            address2 += typeSize * header2[i + 1] * cutting[i][0];
            elemCount *= cutting[i][1] - cutting[i][0] + 1;
        }
        if (rank > 0) {
            address1 += typeSize * cutting[rank - 1][0];
            address2 += typeSize * cutting[rank - 1][0];
            elemCount *= cutting[rank - 1][1] - cutting[rank - 1][0] + 1;
        }

        dvmh_log(TRACE, "copying (peer) %ld bytes (%ld elements)", elemCount * typeSize, elemCount);
        assertCuda(cudaMemcpyPeer(address2, device2, address1, device1, elemCount * typeSize));
    } else {
        // Oh no, we can not copy it by one command... :(
        long curVal;
        for (curVal = cutting[skip][0]; curVal <= cutting[skip][1]; curVal++) {
            long offset1 = offs1 + header1[skip + 1] * curVal;
            long offset2 = offs2 + header2[skip + 1] * curVal;
            perform_copy_peer_internal(rank, typeSize, device1, header1, device2, header2, cutting, matched, skip + 1, offset1, offset2);
        }
    }
}

static int getMatched(int rank, long *header1, long *header2, Interval *cutting) {
    int i, matched = 0;
    long prevSize = 1;
    for (i = rank; i > 1 && header1[i - 1] == header2[i - 1] && header1[i - 1] == prevSize * (cutting[i - 1][1] - cutting[i - 1][0] + 1); i--) {
        prevSize *= cutting[i - 1][1] - cutting[i - 1][0] + 1;
        matched++;
    }
    return matched;
}

static void performCopyCutting(int rank, int typeSize, long *header1, long *header2, Interval *cutting, enum cudaMemcpyKind where) {
    perform_copy_cutting_internal(rank, typeSize, header1, header2, cutting, where, getMatched(rank, header1, header2, cutting), 0, 0, 0);
}

static void performCopyCudaPeer(int rank, int typeSize, int device1, long *header1, int device2, long *header2, Interval *cutting) {
    perform_copy_peer_internal(rank, typeSize, device1, header1, device2, header2, cutting, getMatched(rank, header1, header2, cutting), 0, 0, 0);
}

void performCopyCuda(int rank, int typeSize, int device1, long *header1, int device2, long *header2, Interval *cutting) {
    dvmh_log(TRACE, "performCopyCuda");
    if (device1 >= 0 || device2 >= 0) {
        performCopyCutting(rank, typeSize, header1, header2, cutting, cudaMemcpyDefault);
/*    if (device1 < 0 && device2 >= 0) {
        assertCuda(cudaSetDevice(device2));
        performCopyCutting(rank, typeSize, header1, header2, cutting, cudaMemcpyHostToDevice);
    } else if (device1 >= 0 && device2 < 0) {
        assertCuda(cudaSetDevice(device1));
        performCopyCutting(rank, typeSize, header1, header2, cutting, cudaMemcpyDeviceToHost);
    } else if (device1 >= 0 && device2 >= 0) {
        assertCuda(cudaSetDevice(device1));
#ifdef DONT_USE_CUDA_PEER
        long bytes = typeSize;
        int i;
        for (i = 0; i < rank; i++)
            bytes *= cutting[i][1] - cutting[i][0] + 1;
        char *buf = (char *)malloc(bytes);
        long tmpHeader[rank + 3];
        fillHeader(rank, typeSize, buf, cutting, tmpHeader);
        performCopyCutting(rank, typeSize, header1, tmpHeader, cutting, cudaMemcpyDeviceToHost);
        assertCuda(cudaSetDevice(device2));
        performCopyCutting(rank, typeSize, tmpHeader, header2, cutting, cudaMemcpyHostToDevice);
        free(buf);
#else
        performCopyCudaPeer(rank, typeSize, device1, header1, device2, header2, cutting);
#endif*/
    } else {
        dvmh_log(ERROR, "Copying host to host by CUDA function???");
        exit(1);
    }
}
